#!/bin/bash

# 默认限制比例
LIMIT_PERCENT=${1:-50}
EXCLUDE_PROCESSES="systemd|kthreadd|ksoftirqd|rcu_sched|kyservice|kylinsec|ksm|migration"

# 检测高CPU进程
HIGH_CPU_PROC=$(ps -eo pid,pcpu,comm --sort=-pcpu --no-headers | 
                awk -v exclude="$EXCLUDE_PROCESSES" '$2 > 10 && $0 !~ exclude {print $0; exit}')

if [ -z "$HIGH_CPU_PROC" ]; then
    echo "未发现高CPU进程"
    exit 0
fi

PID=$(echo $HIGH_CPU_PROC | awk '{print $1}')
CPU_USAGE=$(echo $HIGH_CPU_PROC | awk '{print $2}')
PROC_NAME=$(echo $HIGH_CPU_PROC | awk '{print $3}')

echo "限制进程: $PROC_NAME (PID:$PID)"
echo "当前CPU: ${CPU_USAGE}% -> 限制为: ${LIMIT_PERCENT}%"

# 创建cgroup配置
CGROUP_NAME="cpu_limit_${PID}"
CGROUP_DIR="/sys/fs/cgroup/cpu/${CGROUP_NAME}"

# 检查cgroup挂载点
if [ ! -d "/sys/fs/cgroup/cpu" ]; then
    echo "错误：cgroup文件系统未挂载，请检查系统配置"
    exit 1
fi

# 创建cgroup目录
sudo mkdir -p ${CGROUP_DIR} 2>/dev/null
if [ $? -ne 0 ]; then
    echo "错误：创建cgroup失败，请检查权限"
    exit 1
fi

# 设置CPU配额（基于100% = 100000us）
CPU_QUOTA=$(( $LIMIT_PERCENT * 1000 ))  # 转换为毫秒级精度
echo ${CPU_QUOTA} | sudo tee ${CGROUP_DIR}/cpu.cfs_quota_us >/dev/null
echo 100000 | sudo tee ${CGROUP_DIR}/cpu.cfs_period_us >/dev/null

# 添加进程到cgroup
echo ${PID} | sudo tee ${CGROUP_DIR}/cgroup.procs >/dev/null

# 获取进程树的所有PID（防止子进程逃逸限制）
PIDS=$(pstree -p ${PID} | grep -oP '\(\K\d+')
for p in ${PIDS}; do
    echo ${p} | sudo tee ${CGROUP_DIR}/cgroup.procs >/dev/null 2>&1
done

echo "已成功限制进程"
echo "终止限制方法:"
echo "1. 停止进程: sudo kill -9 ${PID}"
echo "2. 删除cgroup: sudo rmdir ${CGROUP_DIR}"